// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/core/plan_node.h>

#include <pollux/substrait/substrait_extension_collector.h>
#include <pollux/substrait/pollux_to_substrait_type.h>
#include <pollux/substrait/proto/substrait/algebra.pb.h>
#include <pollux/vector/constant_vector.h>

namespace kumo::pollux::substrait {
    class PolluxToSubstraitExprConvertor {
    public:
        explicit PolluxToSubstraitExprConvertor(
            const SubstraitExtensionCollectorPtr &extensionCollector)
            : extensionCollector_(extensionCollector) {
        }

        /// Convert Pollux Expression to Substrait Expression.
        /// @param arena Arena to use for allocating Substrait plan objects.
        /// @param expr Pollux expression needed to be converted.
        /// @param inputType The input row Type of the current processed node,
        /// which also equals the output row type of the previous node of the current.
        /// @return A pointer to Substrait expression object allocated on the arena
        /// and representing the input Pollux expression.
        const ::substrait::Expression &toSubstraitExpr(
            google::protobuf::Arena &arena,
            const core::TypedExprPtr &expr,
            const RowTypePtr &inputType);

        /// Convert Pollux Constant Expression to Substrait
        /// Literal Expression.
        /// @param arena Arena to use for allocating Substrait plan objects.
        /// @param constExpr Pollux Constant expression needed to be converted.
        /// @param litValue The Struct that returned literal expression belong to.
        /// @return A pointer to Substrait Literal expression object allocated on
        /// the arena and representing the input Pollux Constant expression.
        const ::substrait::Expression_Literal &toSubstraitExpr(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::ConstantTypedExpr> &constExpr,
            ::substrait::Expression_Literal_Struct *litValue = nullptr);

        /// Convert Pollux FieldAccessTypedExpr to Substrait FieldReference Expression.
        const ::substrait::Expression_FieldReference &toSubstraitExpr(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::FieldAccessTypedExpr> &fieldExpr,
            const RowTypePtr &inputType);

        /// Convert Pollux vector to Substrait literal.
        const ::substrait::Expression_Literal &toSubstraitLiteral(
            google::protobuf::Arena &arena,
            const pollux::VectorPtr &vectorValue,
            ::substrait::Expression_Literal_Struct *litValue);

    private:
        /// Convert Pollux Cast Expression to Substrait Cast Expression.
        const ::substrait::Expression_Cast &toSubstraitExpr(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::CastTypedExpr> &castExpr,
            const RowTypePtr &inputType);

        /// Convert Pollux CallTypedExpr Expression to Substrait Expression.
        const ::substrait::Expression &toSubstraitExpr(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::CallTypedExpr> &callTypeExpr,
            const RowTypePtr &inputType);

        /// Convert Pollux variant to Substrait Literal Expression.
        const ::substrait::Expression_Literal &toSubstraitLiteral(
            google::protobuf::Arena &arena,
            const pollux::variant &variantValue);

        /// Convert values in Pollux array vector to Substrait Literal List.
        const ::substrait::Expression_Literal_List &toSubstraitLiteralList(
            google::protobuf::Arena &arena,
            const ArrayVector *array_vector,
            vector_size_t row);

        /// Convert values in Pollux complex vector to Substrait Literal.
        const ::substrait::Expression_Literal &toSubstraitLiteralComplex(
            google::protobuf::Arena &arena,
            const std::shared_ptr<ConstantVector<ComplexType> > &constant_vector);

        PolluxToSubstraitTypeConvertorPtr typeConvertor_;

        SubstraitExtensionCollectorPtr extensionCollector_;
    };

    using PolluxToSubstraitExprConvertorPtr =
    std::shared_ptr<PolluxToSubstraitExprConvertor>;
} // namespace kumo::pollux::substrait
